#!/usr/bin/env python

"""Load PDB, add surface profiles and related data"""

import numpy as N
import glob

import Biskit.tools as T
from Biskit import PDBModel, PDBDope

def addSideChainExposure( m ):
    """
    Convert per-atom accessibility into side chain accessibility per residue
    """
    mask_sc = N.logical_not( m.maskBB() ) ## all side chain atoms

    ## mark residues with any side chain atom (all but GLY)
    resmask = m.atom2resMask( mask_sc )

    ## average relAS over all side chain atoms of each residue
    m_sc = m.compress( mask_sc )  ## model with only side chain atoms, lacks GLY
    scASA= m_sc.atom2resProfile( 'relAS', N.average ) ## now per residue

    ## put side chain average into residue profile, 0 for Glycines
    m.residues.set( 'scAS', scASA, mask=resmask, default=0 )


def addResidueExposure( m ):
    """Convert relative per-atom accessibility into per-residue accesssibility."""
    m.residues['resAS'] = m.atom2resProfile( 'relAS', N.average )
    

def exposedO( m ):
    """
    Mark exposed terminal oxygens of SER, THR, TYR 
    """
    m_o = m.mask( ['OG', 'OG1', 'OH'] )
    return N.logical_and( m['relAS'] > 10., m_o )
    

def reportPdb( m, fname ):
    """
    create PDB file with side chain exposure mapped to B-Factor column and
    exposed S/T/Y Oxygen labelled in occupancy column.
    """
    m['temperature_factor'] = m.res2atomProfile('scAS')
    m['occupancy'] = N.logical_not(m['exo']) * 1. ## 0 means exposed O

    m.writePdb( fname )


def tabify( l ):
    """convert list or array to tab-delimited string."""
    return '\t'.join( [str(x) for x in l] )

def tabifyRounded( l ):
    """convert list or array to tab-delimited string."""
    return '\t'.join( [str(round(x,2)) for x in l] )


def reportExposedSTY( m, fname ):
    """Report exposed S/T/Y in a text file"""
    try:
        f = open(fname, 'w')
        f.write(
            '# (1) PDB sequence, (2) exposed S/T/Y (position 0...N), (3) all S/T/Y\n')
        
        f.write( m.sequence() + '\n' )

        mask_exo = m.atom2resProfile( m['exo'], N.sum )
        f.write( tabify( N.flatnonzero( mask_exo ) )  + '\n')

        mask_sty = N.array( [ s in ['S','T','Y'] for s in m.sequence() ] )
        f.write( tabify( N.flatnonzero( mask_sty ) ) + '\n')
        
    except IOError, why:
        print 'error creating report file %s:\n %r' % (fname, why) 


def reportSurface( m, fname ):
    """Report surface profile and model details in a text file"""
    try:
        f = open( fname, 'w' )
        f.write('# lines: (1) PDB seq.; (2) side-chain expos.; (3) total expos.; (4) PDB start TAB end TAB template TAB sequence_id; (4) STY positions\n')

        f.write( m.sequence() + '\n' )

        f.write( tabifyRounded( m['scAS'] )  + '\n')
        f.write( tabifyRounded( m['resAS'] ) + '\n' )

        f.write(
            'MODEL: %(target_begin)i\t%(target_end)i\t%(template)s\t%(seq_identity)5.2f\n' \
            % m.info )

        mask_sty = N.array( [ s in ['S','T','Y'] for s in m.sequence() ] )
        f.write( 'STY: ' + tabify( N.flatnonzero( mask_sty ) ) + '\n')
        
    except IOError, why:
        print 'error creating report file %s:\n %r' % (fname, why) 

        

def prepareModel( fname ):
    """Load PDB and add surface and other profiles"""

    patterns = [ ('seq_identity', 'SEQUENCE IDENTITY: *([0-9\.]+)'),
                 ('model_score',  'MODEL SCORE: *([0-9\.]+)'),
                 ('quality_score', 'ModPipe Quality Score: *([0-9\.]+)'),
                 ('e_value', 'EVALUE: *([0-9\.]+)'),
                 ('z_dope', 'zDOPE: *([\-0-9\.]+)'),
                 ('template', 'TEMPLATE PDB: *([0-9a-zA-Z]+)'),
                 ('template_chain', 'TEMPLATE CHAIN: *([A-Z]+)'),
                 ('target_length', 'TARGET LENGTH: *([0-9\.]+)'),
                 ('target_begin', 'TARGET BEGIN: *([0-9\.]+)'),
                 ('target_end', 'TARGET END: *([0-9\.]+)'),
                 ('template_begin', 'TEMPLATE BEGIN: *([0-9\.]+)'),
                 ('template_end', 'TEMPLATE END: *([0-9\.]+)'),
                 ('modpipe_id', 'MODPIPE MODEL ID: *([0-9a-z]+)'),
                 ('modpipe_aln', 'MODPIPE ALIGNMENT ID: *([0-9a-z]+)') ]
                 

    m = PDBModel( fname, headPatterns=patterns )

    m['modscore'] = m['temperature_factor']  ## copy quality score, perhaps useful
    m['aa'] = m.sequence() ## create residue profile with single-letter AA 

    d = PDBDope( m )
    d.addSurfaceRacer()  ## accessibilities are now in atom profile m['relAS']

    addResidueExposure( m )    ## relative residue exposure now in residue profile 'resAS'
    addSideChainExposure( m )  ## side chain exposure now in residue profile 'scAS'
    m['exo'] = exposedO( m )   ## exposed T/Y/S oxygen now marked in atom profile 'exo'

    return m
        
    
####### MAIN ############

## fname = 'Q99383.6.pdb'

fnames = glob.glob( '*.pdb' )

for f in fnames:

    T.flushPrint( 'processing %s...' % f )
    try:
        m = prepareModel( f )
        m.saveAs( T.stripFilename( f ) + '.model' ) ## save object to disc

        reportSurface( m, T.stripFilename( f ) + '_surface.txt' )

        print 'OK'

    except:
        print 'Error'
        T.errWriteln( T.lastError() )

##     reportPdb(     m, T.stripFilename( fname ) + '_exposure.pdb' )
